import cv2
import os
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional

class ImagePreprocessor:
    def __init__(self, debug_mode: bool = False):
        self.debug = debug_mode
    
    # 第一次预处理
    def process_grid_image(self, image_path):
        # 1. 加载图片
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法读取图片，请检查路径")
        
        # 2. 识别背景干扰色（左下角25x90区域）
        bg_area = img[210:300, 0:25]  # 左下角区域
        bg_color = np.median(bg_area, axis=(0,1))  # 背景色中值
        
        # 3. 去除背景干扰
        dist = np.linalg.norm(img - bg_color, axis=2)
        # mask = np.where(dist < 50, 0, 255).astype(np.uint8)  # 距离阈值
        mask = np.where(dist < 20, 0, 255).astype(np.uint8)
        
        # 4. 清除背景
        result = cv2.bitwise_and(img, img, mask=mask)
        result[mask==0] = 255  # 背景设为白色  
        return result
    # 第二次预处理  
    def _show_stage(self, title: str, img: np.ndarray):
        """调试模式下显示处理阶段结果"""
        if self.debug:
            if len(img.shape) == 2:  # 灰度图
                plt.imshow(img, cmap='gray')
            else:
                plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            plt.title(title)
            plt.axis('off')
            plt.show()

    def advanced_denoise(self, img: np.ndarray) -> np.ndarray:
        """混合降噪策略"""
        try:
            # 彩色降噪
            denoised = cv2.fastNlMeansDenoisingColored(img, None, 
                      h=15, hColor=13, 
                      templateWindowSize=7, 
                      searchWindowSize=21)
            
            # 灰度降噪二次处理
            gray = cv2.cvtColor(denoised, cv2.COLOR_BGR2GRAY)
            denoised_gray = cv2.fastNlMeansDenoising(gray, None, 20, 7, 21)
            
            self._show_stage("Denoised", denoised_gray)
            return denoised_gray
        except Exception as e:
            print(f"Denoising error: {str(e)}")
            return img

    def remove_lines(self, img: np.ndarray) -> np.ndarray:
        """消除干扰线条"""
        try:
            # 多尺度线条检测
            edges = cv2.Canny(img, 50, 150, apertureSize=3)
            lines = cv2.HoughLinesP(edges, 1, np.pi/180, 
                    threshold=100, 
                    minLineLength=100, 
                    maxLineGap=10)
            
            if lines is not None:
                mask = np.ones_like(img) * 255
                for line in lines:
                    x1, y1, x2, y2 = line[0]
                    thickness = max(3, int(abs(x2-x1)/50))  # 动态线条粗细
                    cv2.line(mask, (x1,y1), (x2,y2), (0,0,0), thickness)
                
                # 改进的图像修复
                result = cv2.inpaint(img, 255-mask[:,:,0], 
                           inpaintRadius=3, 
                           flags=cv2.INPAINT_TELEA)
            else:
                result = img.copy()
                
            self._show_stage("Lines Removed", result)
            return result
        except Exception as e:
            print(f"Line removal error: {str(e)}")
            return img

    def contrast_enhance(self, img: np.ndarray) -> np.ndarray:
        """自适应对比度增强"""
        try:
            # 动态CLAHE参数
            clahe = cv2.createCLAHE(
                clipLimit=2.0 + img.std()/25,  # 基于图像噪声动态调整
                tileGridSize=(8,8))
            
            enhanced = clahe.apply(img) if len(img.shape)==2 else \
                      clahe.apply(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
            
            # 智能伽马校正
            hist = cv2.calcHist([enhanced],[0],None,[256],[0,256])
            cdf = hist.cumsum()
            cdf_normalized = cdf * hist.max() / cdf.max()
            gamma = 0.4 + 1.5 * (1 - cdf_normalized[150]/255)  # 动态gamma
            
            lookup = np.array([((i / 255.0) ** (1/gamma)) * 255 
                      for i in np.arange(0, 256)]).astype("uint8")
            result = cv2.LUT(enhanced, lookup)
            
            self._show_stage("Contrast Enhanced", result)
            return result
        except Exception as e:
            print(f"Contrast enhancement error: {str(e)}")
            return img

    def adaptive_binarization(self, img: np.ndarray) -> np.ndarray:
        """智能二值化"""
        try:
            # 大津法全局阈值
            _, global_thresh = cv2.threshold(img, 0, 255, 
                      cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            # 改进的自适应阈值
            block_size = min(31, max(11, img.shape[1]//30))  # 动态块大小
            adaptive = cv2.adaptiveThreshold(img, 255, 
                      cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                      cv2.THRESH_BINARY, block_size, 5)
            
            # 混合策略
            result = cv2.bitwise_and(global_thresh, adaptive)
            
            # 后处理
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2,2))
            result = cv2.morphologyEx(result, cv2.MORPH_CLOSE, kernel)
            
            self._show_stage("Binarized", result)
            return result
        except Exception as e:
            print(f"Binarization error: {str(e)}")
            return img

    def full_pipeline(self, img_path: str) -> Optional[np.ndarray]:
        """完整处理流程"""
        try:
            # 1. 读取图像
            img = cv2.imread(img_path)
            if img is None:
                raise ValueError("Image not found or corrupted")
            
            # 2. 降噪处理
            denoised = self.advanced_denoise(img)
            
            # 3. 去除干扰线
            no_lines = self.remove_lines(denoised)
            
            # 4. 对比度增强
            enhanced = self.contrast_enhance(no_lines)
            
            # 5. 二值化
            binary = self.adaptive_binarization(enhanced)
            
            # 6. 显示最终结果
            if self.debug:
                self._show_stage("Final Result", binary)
            
            return binary
        except Exception as e:
            print(f"Pipeline error: {str(e)}")
            return None


# 使用示例
if __name__ == "__main__":
    # 初始化处理器（开启调试模式）
    processor = ImagePreprocessor(debug_mode=False)
    os.chdir(r'E:\Python\vscode\Crawlers\yuanrenxue\no8')
    image_path = r"E:\Python\vscode\Crawlers\yuanrenxue\no8\org_image.jpg"
    target_chars = ['孟', '擋', '驴', '妩']
    
    # 处理图片
    clean_img = processor.process_grid_image(image_path)
    # 保存处理结果
    cv2.imwrite("clean_background.jpg", clean_img)
    # 运行第二次处理
    result = processor.full_pipeline(r"E:\Python\vscode\Crawlers\yuanrenxue\no8\clean_background.jpg") 
    # 保存处理结果
    if result is not None:
        cv2.imwrite("processed_result.png", result)

